import matplotlib.pyplot as plt
import torch
import numpy as np

# (1)数据处理
# ①读取data-02-stock_daily.csv
data = np.loadtxt('data-02-stock_daily.csv', delimiter=',')
data = data[::-1]
# ②按照要求处理x，y
# ③归一化数值
from sklearn.preprocessing import MinMaxScaler

data = MinMaxScaler().fit_transform(data)
c = 7
x = []
y = []
for i in range(len(data) - c):
    x_data = data[i:i + c]
    y_data = data[i + c][-1]
    x.append(x_data)
    y.append(y_data)
x = torch.Tensor(x)
y = torch.Tensor(y).reshape(-1, 1)
print()
# ④数据切分处理
from sklearn.model_selection import train_test_split

train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)
print(train_x.shape)
# ②创建线性回归模型
model = torch.nn.Linear(in_features=35, out_features=1)
# ③设置mse损失
loss_fn = torch.nn.MSELoss()
# ④设置随机梯度下降，学习率自定
op = torch.optim.Adam(params=model.parameters(), lr=0.001)
# (2)模型处理
# ①循环2000次
loss_list = []
for i in range(2000):
    # ②初始化梯度
    op.zero_grad()
    h = model(train_x.reshape(-1, 35))
    loss = loss_fn(h, train_y)
    loss_list.append(loss)
    if i % 10 == 0:
        print(loss)
    loss.backward()
    # ③进行梯度下降处理
    op.step()

plt.plot(test_y, c='r')
plt.plot(model(test_x.reshape(-1, 35)).data.numpy(), c='g')
plt.show()
# ④每10次打印损失值
# ⑤打印预测结果
# ⑥绘制损失值曲线
